#include <vector>
#include <iostream>
#include <fstream>
#include <complex>
#include <time.h>

#include "ocl.h"
#include "Utils.h"

using namespace std;

typedef complex<float> Complex;

class Lattice {
public:
	vector<int> dim;
	unsigned long volume(int mu=-1) {
		long s=1;
		for(int k=dim.size()-1; k>mu; k--) s*=dim[k];
		return s;
	}
};


class Site {
public:
	Lattice *ptr_lattice;
	unsigned long idx;
	Site(Lattice &lattice, long idx=0) {
		ptr_lattice = &lattice;
		this->idx = idx;
	}
	void set(int x0=0, int x1=0, int x2=0, int x3=0, int x4=0, 
		int x5=0, int x6=0, int x7=0, int x8=0, int x9=0) {
			int n = ptr_lattice->dim.size();
			idx = 0;
			if(n>0) idx += ptr_lattice->volume(0)*x0;
			if(n>1) idx += ptr_lattice->volume(1)*x1;
			if(n>2) idx += ptr_lattice->volume(2)*x2;
			if(n>3) idx += ptr_lattice->volume(3)*x3;
			if(n>4) idx += ptr_lattice->volume(4)*x4;
			if(n>5) idx += ptr_lattice->volume(5)*x5;
			if(n>6) idx += ptr_lattice->volume(6)*x6;
			if(n>7) idx += ptr_lattice->volume(7)*x7;
			if(n>8) idx += ptr_lattice->volume(8)*x8;
			if(n>9) idx += ptr_lattice->volume(9)*x9;
	}
	int operator()(int mu) {
		return (idx / ptr_lattice->volume(mu)) % ptr_lattice->dim[mu];
	}
};

Site operator+(const Site &x, unsigned int mu) {
	// shift right on the lattice of one step in direction mu=+0,+1,+2,...
	mu = mu % x.ptr_lattice->dim.size();
	long idx = x.idx;
	long v = x.ptr_lattice->volume(mu);
	int n = x.ptr_lattice->dim[mu];
	long projection = (idx / v) % n;
	idx += ((projection+1) % x.ptr_lattice->dim[mu] - projection) * v;
	Site y(*x.ptr_lattice,idx);  
	return y;
}

Site operator-(const Site &x, unsigned int mu) {
	// shift left on the lattice of one step in direction mu=-0,-1,-2,...
	mu = mu % x.ptr_lattice->dim.size();
	long idx = x.idx;
	long v = x.ptr_lattice->volume(mu);
	int n = x.ptr_lattice->dim[mu];
	long projection = (idx / v) % n;
	idx += ((projection+n-1) % n - projection) * v;
	Site y(*x.ptr_lattice,idx);  
	return y;
}

class Gauge {
public:
	Lattice *ptr_lattice;
	int nc;
	vector<Complex> m;
	Gauge(Lattice &lattice, int nc) {
		ptr_lattice = &lattice;
		this->nc=nc;
		m.resize(lattice.volume()*lattice.dim.size()*nc*nc);
	}
	int ndim() const {
		return ptr_lattice->dim.size();
	}
	const Complex get(const Site &x, int mu, int i, int j) const {
		int n = ndim(); 
		if(mu>=0) return m[((x.idx*n + (mu % n))*nc+i)*nc+j];
		else return conj(m[((x.idx*n + (-mu % n))*nc+j)*nc+i]);
	}
	Complex &operator()(const Site &x, int mu, int i, int j) {
		int n = ndim(); 
		if(mu<0) throw "error";
		return m[((x.idx*ndim() + (mu % n))*nc+i)*nc+j];
	}
};

class PathProduct {
public:
	Lattice *ptr_lattice;
	int nc;
	vector<Complex> m;
	PathProduct(Lattice &lattice,int nc) {
		ptr_lattice = &lattice;
		this->nc=nc;
		m.resize(lattice.volume()*nc*nc);
	}
	const Complex get(const Site &x, int i, int j) const {
		return m[(x.idx*nc+i)*nc+j];
	}
	Complex &operator()(const Site &x, int i, int j) {
		return m[(x.idx*nc+i)*nc+j];
	}
};

typedef vector<int> Path;

void performFirstLoop(OpenCL &cl, Gauge &U, Lattice &spacetime, int nc) {
	try {
		
		vector<float> out;
		out.resize(U.m.size());
		int sz = spacetime.dim.size();
		cl.setArgumentReadOnly<int *>(0, sizeof(int), &sz);
		cl.setArgumentReadOnly<int *>(1, sizeof(int), &nc);
		cl.setArgumentWriteOnly<float *>(2, sizeof(float) * U.m.size(), &out[0]);
		
		//Run the kernel
		cl.execute(spacetime.volume());
		
		//Read back the data
		cl.readArgument(2, sizeof(float) * U.m.size(), &out[0]);
		for(unsigned int i = 0; i < out.size(); i++) {
			U.m[i] = out[i];
		}
	} catch ( std::string ex ) {
		cout << ex << endl;
		cout.flush();
	}
}

void performSecondLoop(OpenCL &loop2, Site &x, Site &y, Lattice &spacetime, Path &path, PathProduct &P, Gauge &U, int nc) {
	float *tmp = new float[nc * nc];
	for(int i = 0; i < (nc * nc); i++) {
		tmp[i] = 0;
	}

	vector<float> path_m;
	path_m.resize(P.m.size());
	for(unsigned int i = 0; i < path_m.size(); i++) {
		path_m[i] = P.m[i].real();
	}

	vector<float> gauge_m;
	gauge_m.resize(U.m.size());
	for(unsigned int i = 0; i < gauge_m.size(); i++) {
		gauge_m[i] = U.m[i].real();
	}

	try {

		int sz = spacetime.dim.size();
		int path_size = path.size();
		loop2.setArgumentReadOnly<int *>(0, sizeof(int), &sz);
		loop2.setArgumentReadOnly<float *>(1, sizeof(float) * gauge_m.size(), &gauge_m[0]);
		loop2.setArgumentReadOnly<int *>(2, sizeof(int), &path_size);
		loop2.setArgumentReadOnly<int *>(3, sizeof(int) * path.size(), &path[0]);
		loop2.setArgumentWriteOnly<float *>(4, sizeof(float) * path_m.size(), &path_m[0]);
		loop2.setArgumentReadOnly<int *>(5, sizeof(int), &nc);
		loop2.setArgumentWriteOnly<float *>(6, sizeof(float) * (nc * nc), tmp);
		loop2.execute(spacetime.volume());
		loop2.readArgument<float *>(4, sizeof(float) * path_m.size(), &path_m[0]);
	} catch ( std::string ex ) {
		cout << ex << endl;	
		cout.flush();
	}

	for(unsigned int i = 0; i < path_m.size(); i++) {
		P.m[i] = path_m[i];
	}

	for(unsigned int i = 0; i < gauge_m.size(); i++) {
		U.m[i] = gauge_m[i];
	}

	delete [] tmp;
}

int orgMain();

int main() {

	clock_t start1 = clock();
	orgMain();
	clock_t end1 = clock();

	// input nc
	int nc=5;
	// input Lattice size
	
	#ifdef _nightst4r_
		//int devType = CL_DEVICE_TYPE_CPU;
		int devType = CL_DEVICE_TYPE_GPU;
	#else 
		int devType = CL_DEVICE_TYPE_GPU;
	#endif

	cout << "Loading OpenCL... " << (devType == CL_DEVICE_TYPE_CPU ? "CPU" : "GPU" ) << endl;
	OpenCL cl("kernel.cl", "input_u", devType);
	cout << cl.ToString() << endl;		

	OpenCL loop2("kernel.cl", "loop2", devType);

	
	clock_t start2 = clock();

	Lattice spacetime;  
	spacetime.dim.resize(4);
	spacetime.dim[0] = spacetime.dim[1] = spacetime.dim[2] = spacetime.dim[3] = 4;
	Site x(spacetime);
	Site y(spacetime);

	// test
	x.set(2,2,2,2);
	//cout << x(0) << x(1) << x(2) << x(3) << endl;
	x = x+2;
	//cout << x(0) << x(1) << x(2) << x(3) << endl;
	x=x-0;
	//cout << x(0) << x(1) << x(2) << x(3) << endl;
	// end test		
	
	// input U
	// defines ndim x nc x nc matrices at each point of spacetime
	Gauge U(spacetime,nc);
	performFirstLoop(cl, U, spacetime, nc);

	// input path
	Path path(6);
	path[0]=+1; path[1]=+1; path[2]=+2; path[3]=-1; path[4]=-1; path[5]=-2;

	// at every point of lattice compute P
	PathProduct P(spacetime,nc);  
	//Complex tmp[nc][nc];
	performSecondLoop(loop2, x,y,spacetime,path,P,U,nc);

	clock_t end2 = clock();
	
	cout << endl;
	cout << "SINGLE CPU TIME: " <<  (float)(end1 - start1) / (float)CLOCKS_PER_SEC << "S" << endl;
	cout << "OpenCL TIME: " <<  (float)(end2 - start2) / (float)CLOCKS_PER_SEC << "S" << endl;

	cout << endl;

	for (unsigned int i = 0; i < P.m.size(); i++ ) {
		cout << P.m[i].real() << " " ;
	}
	cout << endl;
	
	
	
	cout << "No errors!  Press any key and Enter to continue" << endl;
	getchar();

	return 0;
}

// NOTES: the loops on x.idx need to be parallelized (the looping order is not relevant).








